import os
import pandas as pd
from sklearn.model_selection import train_test_split
from tqdm import tqdm


def load_thucnews_data(data_dir):
    """
    加载THUCNews数据集，将类别映射为连续整数。

    参数:
        data_dir (str): THUCNews数据集的根目录路径（包含多个子文件夹）。

    返回:
        df (pd.DataFrame): 包含'label'和'text_a'两列的数据框。
        label_map (dict): 原始类别名到整数的映射字典。
    """
    # 获取所有类别文件夹
    categories = [name for name in os.listdir(data_dir) if os.path.isdir(os.path.join(data_dir, name))]
    categories.sort()  # 确保类别顺序固定

    # 创建类别到整数的映射
    label_map = {category: idx for idx, category in enumerate(categories)}

    # 收集所有数据
    data = []
    for category in tqdm(categories, desc="Loading categories"):
        category_dir = os.path.join(data_dir, category)
        for filename in os.listdir(category_dir):
            file_path = os.path.join(category_dir, filename)
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    # 处理多行文本为单行
                    content = ' '.join(f.read().split())  # 关键修改
                    if content:
                        data.append((label_map[category], content))
            except Exception as e:
                print(f"Error reading {file_path}: {e}")


    # 转换为DataFrame
    df = pd.DataFrame(data, columns=['label', 'text_a'])
    return df, label_map


def split_and_save(df, output_dir, test_size=0.1, dev_size=0.1, random_seed=42):
    """
    划分数据集并保存为TSV文件。

    参数:
        df (pd.DataFrame): 需要划分的数据框。
        output_dir (str): 输出目录路径。
        test_size (float): 测试集比例。
        dev_size (float): 验证集比例。
        random_seed (int): 随机种子。
    """
    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"创建目录: {output_dir}")

    # 划分数据集
    train_dev_df, test_df = train_test_split(
        df,
        test_size=test_size,
        random_state=random_seed,
        stratify=df['label']
    )

    dev_size_relative = dev_size / (1 - test_size)
    train_df, dev_df = train_test_split(
        train_dev_df,
        test_size=dev_size_relative,
        random_state=random_seed,
        stratify=train_dev_df['label']
    )

    # 保存为TSV文件
    train_df.to_csv(os.path.join(output_dir, "train.tsv"), sep='\t', index=False, header=True)
    dev_df.to_csv(os.path.join(output_dir, "dev.tsv"), sep='\t', index=False, header=True)
    test_df.to_csv(os.path.join(output_dir, "test.tsv"), sep='\t', index=False, header=True)


# 使用示例
if __name__ == "__main__":
    data_dir = "E:/词向量/dataset/中文多分类/THUCNews"  # 替换为您的THUCNews数据集路径
    output_dir = "./thucnews_tsv"  # 输出目录

    # 加载数据并获取标签映射
    df, label_map = load_thucnews_data(data_dir)
    print("标签映射关系：", label_map)

    # 划分并保存数据集
    split_and_save(df, output_dir, test_size=0.1, dev_size=0.1, random_seed=42)